# coding: utf-8

import demo_common
from gensim import models
import os

model_path = './model_dir'
if not os.path.exists(model_path):
    os.mkdir(model_path)


def get_model_word2vec(train_set, need_retrain=False):
    path = '{}/{}'.format(model_path, 'word2vec')
    if not need_retrain and os.path.exists(path):
        return models.Word2Vec.load(path)
    else:
        print('start training')
        model = models.Word2Vec(sentences=train_set, size=100, window=5, min_count=2, workers=4, sg=1)
        model.save(path)
        return model


if __name__ == '__main__':
    train_set = demo_common.get_train_set(max_count=55000)
    print('train_set len', len(train_set))
    # print(train_set)
    word2vec = get_model_word2vec(train_set=train_set, need_retrain=True)
    # 查看相近的词
    print(word2vec.wv.most_similar(['孩子']))
    # 获取对应的词向量
    print(word2vec.wv['孩子'])
